import cv2
import numpy as np
from pathlib import Path
import os
import time

def create_output_dir(output_dir):
    """创建输出目录"""
    os.makedirs(output_dir, exist_ok=True)
    return output_dir

def calculate_frame_clarity(frame):
    """计算帧的清晰度"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    return cv2.Laplacian(gray, cv2.CV_64F).var()

def detect_scene_change(prev_frame, curr_frame, min_threshold=0.15, max_threshold=0.60):
    """
    检测场景变化
    :param prev_frame: 前一帧
    :param curr_frame: 当前帧
    :param min_threshold: 最小变化阈值
    :param max_threshold: 最大变化阈值
    :return: (是否是新场景，变化程度)
    """
    # 转换为灰度图
    prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
    curr_gray = cv2.cvtColor(curr_frame, cv2.COLOR_BGR2GRAY)
    
    # 计算帧差
    diff = cv2.absdiff(prev_gray, curr_gray)
    # 应用高斯模糊减少噪声
    diff = cv2.GaussianBlur(diff, (5, 5), 0)
    # 应用阈值
    _, thresh = cv2.threshold(diff, 25, 255, cv2.THRESH_BINARY)
    # 计算变化像素的比��
    change_ratio = np.count_nonzero(thresh) / thresh.size
    # 判断是否是新场景
    is_new_scene = min_threshold < change_ratio < max_threshold
    
    return is_new_scene, change_ratio

def split_video_scenes(video_path, output_dir, min_threshold=0.15, max_threshold=0.60, min_scene_duration=15):
    """
    分割视频场景
    :param video_path: 视频文件路径
    :param output_dir: 输出目录
    :param min_threshold: 最小变化值
    :param max_threshold: 最大变化阈值
    :param min_scene_duration: 最小场景持续帧数
    :return: 场景图片路径列表
    """
    print(f"开始处理视频: {video_path}")
    start_time = time.time()
    
    # 创建输出目录
    output_dir = create_output_dir(output_dir)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("无法打开视频文件")
    
    # 获取视频信息
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    print(f"视频信息: {total_frames}帧，{fps} FPS")
    
    scene_images = []
    scene_count = 0
    frame_count = 0
    frames_since_last_scene = 0
    prev_frame = None
    current_scene_frames = []
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        # 显示进度
        if frame_count % 100 == 0:
            progress = (frame_count / total_frames) * 100
            print(f"处理进度: {progress:.1f}%")
            
        # 第一帧处理
        if prev_frame is None:
            prev_frame = frame
            current_scene_frames.append((frame, calculate_frame_clarity(frame)))
            continue
            
        # 检测场景变化
        is_new_scene, change_ratio = detect_scene_change(prev_frame, frame, min_threshold, max_threshold)
        
        # 将当前帧添加到当前场景
        current_scene_frames.append((frame, calculate_frame_clarity(frame)))
        frames_since_last_scene += 1
        
        # 如果检测到新场景且已经过了最小场景持续时间
        if is_new_scene and frames_since_last_scene >= min_scene_duration:
            # 从当前场景帧中选择最清晰的帧
            clearest_frame, _ = max(current_scene_frames, key=lambda x: x[1])
            
            # 保存场景图片
            scene_path = os.path.join(output_dir, f'scene_{scene_count:03d}.jpg')
            cv2.imwrite(scene_path, clearest_frame)
            scene_images.append(str(Path(scene_path).absolute()))
            
            print(f"检测到新场景 {scene_count}，在帧 {frame_count}，变化程度: {change_ratio:.3f}")
            
            # 重置计数器和缓存
            scene_count += 1
            frames_since_last_scene = 0
            current_scene_frames = [(frame, calculate_frame_clarity(frame))]
            
        prev_frame = frame
        frame_count += 1
        
    # 保存最后一个场景
    if current_scene_frames:
        clearest_frame, _ = max(current_scene_frames, key=lambda x: x[1])
        scene_path = os.path.join(output_dir, f'scene_{scene_count:03d}.jpg')
        cv2.imwrite(scene_path, clearest_frame)
        scene_images.append(str(Path(scene_path).absolute()))
    
    # 释放资源
    cap.release()
    
    # 打印统计信息
    end_time = time.time()
    processing_time = end_time - start_time
    print(f"\n处理完成:")
    print(f"总共检测到 {scene_count + 1} 个场景")
    print(f"处理时间: {processing_time:.2f} 秒")
    print(f"平均每秒处理 {frame_count/processing_time:.1f} 帧")
    
    return scene_images

def main():
    # 配置参数
    video_path = "downloads/a.mp4"
    output_dir = "output_scenes"
    
    # 添加文件检查
    if not os.path.exists(video_path):
        print(f"错误：视频文件不存在: {video_path}")
        return
        
    try:
        scene_paths = split_video_scenes(
            video_path=video_path,
            output_dir=output_dir,
            min_threshold=0.15,
            max_threshold=0.60,
            min_scene_duration=15
        )
        
        print("\n场景图片路径:")
        for path in scene_paths:
            print(path)
    except Exception as e:
        print(f"处理视频时出错: {str(e)}")

if __name__ == "__main__":
    main()
